package simpledb.optimizer;

import simpledb.ParsingException;
import simpledb.common.Database;
import simpledb.execution.*;

import javax.swing.*;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeCellRenderer;
import java.util.*;

/**
 * JoinOptimizer类负责以最优方式排序一系列连接操作，
 * 并为给定的逻辑计划选择最佳的连接实例化方式。
 */
public class JoinOptimizer {

    final LogicalPlan           p;    // 逻辑计划
    final List<LogicalJoinNode> joins; // 要执行的连接列表

    /**
     * 构造函数
     *
     * @param p 优化中的逻辑计划
     * @param joins 要执行的连接列表
     */
    public JoinOptimizer(LogicalPlan p, List<LogicalJoinNode> joins) {
        this.p = p;
        this.joins = joins;
    }

    /**
     * 返回给定逻辑连接的最佳迭代器，根据指定的统计信息和左右子计划。
     * 注意：由于OpIterator不提供任何基数估计，而stats只有关于基表的信息，
     * 因此缺少确定哪个计划应为内部/外部计划的信息。
     *
     * @param lj 正在考虑的连接
     * @param plan1 左连接节点的子节点
     * @param plan2 右连接节点的子节点
     */
    public static OpIterator instantiateJoin(LogicalJoinNode lj, OpIterator plan1, OpIterator plan2)
                                                                                                    throws ParsingException {

        int t1id = 0, t2id = 0;
        OpIterator j;

        try {
            // 获取左计划中字段的索引
            t1id = plan1.getTupleDesc().fieldNameToIndex(lj.f1QuantifiedName);
        } catch (NoSuchElementException e) {
            throw new ParsingException("未知字段 " + lj.f1QuantifiedName);
        }

        if (lj instanceof LogicalSubplanJoinNode) {
            // 如果是子计划连接，右表ID设为0
            t2id = 0;
        } else {
            try {
                // 获取右计划中字段的索引
                t2id = plan2.getTupleDesc().fieldNameToIndex(lj.f2QuantifiedName);
            } catch (NoSuchElementException e) {
                throw new ParsingException("未知字段 " + lj.f2QuantifiedName);
            }
        }

        // 创建连接谓词
        JoinPredicate p = new JoinPredicate(t1id, lj.p, t2id);

        // 如果是等值连接，尝试使用HashEquiJoin，如果失败则回退到普通Join
        if (lj.p == Predicate.Op.EQUALS) {
            try {
                // 动态加载HashEquiJoin - 如果不存在，则回退到常规连接
                Class<?> c = Class.forName("simpledb.execution.HashEquiJoin");
                java.lang.reflect.Constructor<?> ct = c.getConstructors()[0];
                j = (OpIterator) ct.newInstance(new Object[] { p, plan1, plan2 });
            } catch (Exception e) {
                j = new Join(p, plan1, plan2);
            }
        } else {
            // 非等值连接只能使用普通Join
            j = new Join(p, plan1, plan2);
        }

        return j;
    }

    /**
     * 估计连接的成本。
     *
     * 连接成本应基于Lab 2中实现的连接算法计算。
     * 成本计算基于查询过程中必须读取的数据量以及连接执行的CPU操作数。
     * 假设单个谓词应用的成本大约为1。
     *
     * @param j 表示正在执行的连接操作的LogicalJoinNode
     * @param card1 左侧数据集的估计基数
     * @param card2 右侧数据集的估计基数
     * @param cost1 左侧表完整扫描的估计成本
     * @param cost2 右侧表完整扫描的估计成本
     * @return 查询成本的估计值，基于cost1和cost2
     */
    public double estimateJoinCost(LogicalJoinNode j, int card1, int card2, double cost1, double cost2) {
        if (j instanceof LogicalSubplanJoinNode) {
            // LogicalSubplanJoinNode表示子查询
            // Lab 3中不需要为此实现完整的支持
            return card1 + cost1 + cost2;
        } else {
            // 插入您的代码
            // 提示：如果您实现了比基本嵌套循环连接更复杂的连接算法
            // 可能需要使用变量"j"

            // I/O成本 = 扫描左表的成本 + 左表每行连接右表的成本
            final double IoCost = cost1 + card1 * cost2;
            // CPU成本 = 左右表行数相乘
            final double CpuCost = card1 * card2;
            return IoCost + CpuCost;
        }
    }

    /**
     * 估计一个连接操作（Join）的结果元组数量（即连接的基数）。
     * 连接的基数是指该连接操作产生的元组数量。
     *
     * @param j      一个表示当前连接操作的 LogicalJoinNode 对象
     * @param card1  连接左侧表的基数（即该表中的元组数量）
     * @param card2  连接右侧表的基数
     * @param t1pkey 左侧表是否是主键表（每个连接字段值唯一）
     * @param t2pkey 右侧表是否是主键表
     * @param stats  数据库中各表的统计信息，以表名作为键（不是别名）
     * @return       估计的连接结果的元组数量（基数）
     */
    public int estimateJoinCardinality(LogicalJoinNode j, int card1, int card2, boolean t1pkey, boolean t2pkey,
                                       Map<String, TableStats> stats) {
        if (j instanceof LogicalSubplanJoinNode) {
            // LogicalSubplanJoinNode 表示的是子查询连接操作
            // 在 Lab3 中不需要支持子查询连接的准确基数估计，直接返回左表的基数即可
            return card1;
        } else {
            // 普通的两个表之间的连接，调用更具体的估计方法来计算连接后的基数
            return estimateTableJoinCardinality(j.p, // 连接谓词操作符（如等于、大于等）
                j.t1Alias, // 左表的别名
                j.t2Alias, // 右表的别名
                j.f1PureName, // 左表连接字段名（不带表名前缀）
                j.f2PureName, // 右表连接字段名（不带表名前缀）
                card1, // 左表基数
                card2, // 右表基数
                t1pkey, // 左表是否为主键表
                t2pkey, // 右表是否为主键表
                stats, // 表统计信息
                p.getTableAliasToIdMapping() // 获取表别名到实际表ID的映射
            );
        }
    }

    /**
     * 估计两个表的连接基数。
     */
    public static int estimateTableJoinCardinality(Predicate.Op joinOp, String table1Alias, String table2Alias,
                                                   String field1PureName, String field2PureName, int card1, int card2,
                                                   boolean t1pkey, boolean t2pkey, Map<String, TableStats> stats,
                                                   Map<String, Integer> tableAliasToId) {

        /**
         * 对于等值连接：
         * - 当其中一个属性是主键时，连接产生的元组数不能超过非主键属性的基数
         *
         * 对于非主键的等值连接：
         * - 很难准确预测输出大小，可以采用简单启发式方法（如两个表基数较大的那个）
         *
         * 对于范围扫描：
         * - 同样难以准确预测大小
         * - 可以假设交叉乘积的一个固定比例（如30%）
         * - 一般来说，范围连接的成本应该大于相同大小表的非主键等值连接
         */
        int card = 1;

        // 根据主键情况估算基数
        if (t1pkey && t2pkey) {
            card = Math.min(card1, card2); // 两边都是主键，取较小的
        } else if (!t1pkey && !t2pkey) {
            card = Math.max(card1, card2); // 两边都不是主键，取较大的
        } else {
            // 一边是主键，另一边不是
            card = t1pkey ? card2 : card1;
        }

        // 根据不同的连接操作符调整基数
        switch (joinOp) {
            case EQUALS: {
                // 等值连接保持原基数
                break;
            }
            case NOT_EQUALS: {
                // 不等于连接，基数为总组合减去可能匹配的组合
                card = card1 * card2 - card;
                break;
            }
            default: {
                // 其他操作符（如范围查询），假设为交叉乘积的1/3
                card = card1 * card2 / 3;
            }
        }

        // 确保基数至少为1
        return card <= 0 ? 1 : card;
    }

    /**
     * 辅助方法，枚举给定向量的所有特定大小的子集。
     *
     * @param v 想要获取子集的向量
     * @param size 感兴趣的子集大小
     * @return 所有指定大小的子集集合
     */
    public <T> Set<Set<T>> enumerateSubsets(List<T> v, int size) {
        final Set<Set<T>> result = new HashSet<>();
        final Set<T> path = new HashSet<>();
        subsetDfs(result, path, v, 0, size);
        return result;
    }

    private <T> void subsetDfs(final Set<Set<T>> result, final Set<T> path, final List<T> v, int curIndex,
                               final int size) {
        if (path.size() == size) {
            result.add(new HashSet<>(path));
            return;
        }
        for (int i = curIndex; i < (v.size() - (size - path.size() - 1)); i++) {
            path.add(v.get(i));
            subsetDfs(result, path, v, i + 1, size);
            path.remove(v.get(i));
        }
    }

    /**
     * 在指定表上计算一个逻辑合理且效率较高的连接顺序。
     * 根据PS4的提示，这应该使用动态规划的方法来实现。
     *
     * @param stats 每个参与连接的表的统计信息，按基础表名引用，而不是别名
     * @param filterSelectivities 每个表上的过滤谓词的选择性，按表别名引用（如果没有别名，则使用基础表名）
     * @param explain 表示代码是否应解释其查询计划还是仅仅执行它
     * @return 存储左深连接顺序的LogicalJoinNode列表，表示它们应该被执行的顺序
     * @throws ParsingException 当stats或filterSelectivities缺少连接中的某个表时，
     *                          或者发生其他内部错误时抛出
     */
    public List<LogicalJoinNode> orderJoins(Map<String, TableStats> stats, Map<String, Double> filterSelectivities,
                                            boolean explain) throws ParsingException {
        final PlanCache pc = new PlanCache(); // 初始化计划缓存
        CostCard costCard = null;

        // 遍历所有可能的子集大小（1到joins.size()）
        for (int i = 1; i <= this.joins.size(); i++) {
            // 枚举当前大小的所有子集
            final Set<Set<LogicalJoinNode>> subsets = enumerateSubsets(this.joins, i);

            // 处理每个子集
            for (final Set<LogicalJoinNode> subPlan : subsets) {
                double bestCost = Double.MAX_VALUE;

                // 尝试将每个连接节点作为最后一个连接操作
                for (final LogicalJoinNode removeNode : subPlan) {
                    // 计算这个子计划的成本和基数
                    final CostCard cc = computeCostAndCardOfSubplan(stats, filterSelectivities, removeNode, subPlan,
                        bestCost, pc);

                    // 找到更优的计划
                    if (cc != null && cc.cost < bestCost) {
                        bestCost = cc.cost;
                        costCard = cc;
                    }
                }

                // 如果找到了有效计划，将其添加到缓存中
                if (bestCost != Double.MAX_VALUE) {
                    pc.addPlan(subPlan, bestCost, costCard.card, costCard.plan);
                }
            }
        }

        // 返回最优的连接顺序，如果未找到则返回原始顺序
        return costCard != null ? costCard.plan : joins;
    }

    /**
     * 这是一个辅助方法，用于计算将joinToRemove连接到joinSet中的成本和基数，
     * 前提是已经计算并存储了所有大小为joinSet.size()-1的子计划。
     *
     * @param stats 所有表的统计信息，按表名引用而不是别名
     * @param filterSelectivities 各表上过滤条件的选择性
     * @param joinToRemove 要从joinSet中移除并尝试连接的连接节点
     * @param joinSet 正在考虑的连接集合
     * @param bestCostSoFar 当前joinSet的最优成本（先前调用computeCostAndCardOfSubplan的最小值）
     * @param pc 连接计划的缓存
     * @return 描述成本、基数和最优子计划的CostCard对象
     * @throws ParsingException 当stats、filterSelectivities或pc对象缺少涉及的表时抛出
     */
    @SuppressWarnings("unchecked")
    private CostCard computeCostAndCardOfSubplan(Map<String, TableStats> stats,
                                                 Map<String, Double> filterSelectivities, LogicalJoinNode joinToRemove,
                                                 Set<LogicalJoinNode> joinSet, double bestCostSoFar, PlanCache pc)
                                                                                                                  throws ParsingException {

        LogicalJoinNode j = joinToRemove;

        List<LogicalJoinNode> prevBest;

        // 验证两个表是否存在
        if (this.p.getTableId(j.t1Alias) == null)
            throw new ParsingException("未知表 " + j.t1Alias);
        if (this.p.getTableId(j.t2Alias) == null)
            throw new ParsingException("未知表 " + j.t2Alias);

        // 获取表名和别名信息
        String table1Name = Database.getCatalog().getTableName(this.p.getTableId(j.t1Alias));
        String table2Name = Database.getCatalog().getTableName(this.p.getTableId(j.t2Alias));
        String table1Alias = j.t1Alias;
        String table2Alias = j.t2Alias;

        // 创建新的连接集合，不包含当前要处理的连接
        Set<LogicalJoinNode> news = new HashSet<>(joinSet);
        news.remove(j);

        double t1cost, t2cost; // 表扫描成本
        int t1card, t2card; // 表基数
        boolean leftPkey, rightPkey; // 主键标志

        // 基本情况 - 连接两个基本关系
        if (news.isEmpty()) {
            prevBest = new ArrayList<>();

            // 获取左表的成本和基数
            t1cost = stats.get(table1Name).estimateScanCost();
            t1card = stats.get(table1Name).estimateTableCardinality(filterSelectivities.get(j.t1Alias));
            leftPkey = isPkey(j.t1Alias, j.f1PureName);

            // 获取右表的成本和基数（如果存在）
            if (table2Alias == null) {
                t2cost = 0;
                t2card = 0;
            } else {
                t2cost = stats.get(table2Name).estimateScanCost();
                t2card = stats.get(table2Name).estimateTableCardinality(filterSelectivities.get(j.t2Alias));
            }
            rightPkey = table2Alias != null && isPkey(table2Alias, j.f2PureName);
        } else {
            // news不是空的 - 找到连接j到news的最佳方式
            prevBest = pc.getOrder(news);

            // 如果没有缓存答案（可能是笛卡尔积），返回null
            if (prevBest == null) {
                return null;
            }

            double prevBestCost = pc.getCost(news); // 前一最佳计划的成本
            int bestCard = pc.getCard(news); // 前一最佳计划的基数

            // 检查j的两个表是否已经在之前的计划中
            if (doesJoin(prevBest, table1Alias)) {
                // j.t1已在prevBest中
                t1cost = prevBestCost; // 左边的成本就是prevBest的成本
                t1card = bestCard;
                leftPkey = hasPkey(prevBest); // 检查是否有主键

                // 右边（t2）的信息
                if (j.t2Alias == null) {
                    t2cost = 0;
                    t2card = 0;
                } else {
                    t2cost = stats.get(table2Name).estimateScanCost();
                    t2card = stats.get(table2Name).estimateTableCardinality(filterSelectivities.get(j.t2Alias));
                }
                rightPkey = j.t2Alias != null && isPkey(j.t2Alias, j.f2PureName);

            } else if (doesJoin(prevBest, j.t2Alias)) {
                // j.t2已在prevBest中
                t2cost = prevBestCost; // 右边的成本就是prevBest的成本
                t2card = bestCard;
                rightPkey = hasPkey(prevBest); // 检查是否有主键

                // 左边（t1）的信息
                t1cost = stats.get(table1Name).estimateScanCost();
                t1card = stats.get(table1Name).estimateTableCardinality(filterSelectivities.get(j.t1Alias));
                leftPkey = isPkey(j.t1Alias, j.f1PureName);

            } else {
                // 如果j.t1或j.t2不在prevBest中，不考虑此计划（避免笛卡尔积）
                return null;
            }
        }

        // 计算两种连接方向的成本
        double cost1 = estimateJoinCost(j, t1card, t2card, t1cost, t2cost);
        LogicalJoinNode j2 = j.swapInnerOuter();
        double cost2 = estimateJoinCost(j2, t2card, t1card, t2cost, t1cost);

        // 选择成本较低的连接方向
        if (cost2 < cost1) {
            boolean tmp;
            j = j2;
            cost1 = cost2;
            tmp = rightPkey;
            rightPkey = leftPkey;
            leftPkey = tmp;
        }

        // 如果成本并不优于已知最佳计划，返回null
        if (cost1 >= bestCostSoFar)
            return null;

        // 创建并填充CostCard对象
        CostCard cc = new CostCard();
        cc.card = estimateJoinCardinality(j, t1card, t2card, leftPkey, rightPkey, stats);
        cc.cost = cost1;
        cc.plan = new ArrayList<>(prevBest); // 使用之前的最佳计划
        cc.plan.add(j); // 添加新连接到末尾
        return cc;
    }

    /**
     * 如果指定表在连接列表中返回true，否则返回false
     */
    private boolean doesJoin(List<LogicalJoinNode> joinlist, String table) {
        for (LogicalJoinNode j : joinlist) {
            if (j.t1Alias.equals(table) || (j.t2Alias != null && j.t2Alias.equals(table)))
                return true;
        }
        return false;
    }

    /**
     * 如果字段是指定表的主键返回true，否则返回false
     *
     * @param tableAlias 表在查询中的别名
     * @param field 字段的原始名称
     */
    private boolean isPkey(String tableAlias, String field) {
        int tid1 = p.getTableId(tableAlias);
        String pkey1 = Database.getCatalog().getPrimaryKey(tid1);

        return pkey1.equals(field);
    }

    /**
     * 如果joinlist中的连接之一连接了主键字段则返回true
     */
    private boolean hasPkey(List<LogicalJoinNode> joinlist) {
        for (LogicalJoinNode j : joinlist) {
            if (isPkey(j.t1Alias, j.f1PureName) || (j.t2Alias != null && isPkey(j.t2Alias, j.f2PureName)))
                return true;
        }
        return false;

    }

    /**
     * 辅助函数，显示一个Swing窗口，包含指定连接列表的树状表示。
     * 参见{@link #orderJoins}，当analyze标志为true时可能想要调用此方法。
     *
     * @param js 要可视化的连接计划
     * @param pc 构建最优计划过程中累积的PlanCache
     * @param stats 基表的表统计信息
     * @param selectivities 各表上过滤条件的选择性
     *                      （通过表别名或无别名时的表名来标识表）
     */
    private void printJoins(List<LogicalJoinNode> js, PlanCache pc, Map<String, TableStats> stats,
                            Map<String, Double> selectivities) {

        // 创建并设置JFrame
        JFrame f = new JFrame("连接计划: " + p.getQuery());
        f.setDefaultCloseOperation(WindowConstants.DISPOSE_ON_CLOSE);
        f.setVisible(true);
        f.setSize(300, 500);

        // 创建映射表，用于存储节点
        Map<String, DefaultMutableTreeNode> m = new HashMap<>();

        DefaultMutableTreeNode root = null, treetop = null;
        HashSet<LogicalJoinNode> pathSoFar = new HashSet<>();
        boolean neither;

        // 遍历每个连接节点构建树形结构
        System.out.println(js);
        for (LogicalJoinNode j : js) {
            pathSoFar.add(j);
            System.out.println("当前路径 = " + pathSoFar);

            // 获取两个表的名称
            String table1Name = Database.getCatalog().getTableName(this.p.getTableId(j.t1Alias));
            String table2Name = Database.getCatalog().getTableName(this.p.getTableId(j.t2Alias));

            neither = true;

            // 创建根节点，包含成本和基数信息
            root = new DefaultMutableTreeNode("连接 " + j + " (成本 =" + pc.getCost(pathSoFar) + ", 基数 = "
                                              + pc.getCard(pathSoFar) + ")");

            DefaultMutableTreeNode n = m.get(j.t1Alias);

            // 处理左表节点
            if (n == null) { // 第一次遇到这个表
                n = new DefaultMutableTreeNode(j.t1Alias
                                               + " (成本 = "
                                               + stats.get(table1Name).estimateScanCost()
                                               + ", 基数 = "
                                               + stats.get(table1Name).estimateTableCardinality(
                                                   selectivities.get(j.t1Alias)) + ")");
                root.add(n);
            } else {
                // 添加已存在的左子节点
                root.add(n);
                neither = false;
            }
            m.put(j.t1Alias, root);

            // 处理右表节点
            n = m.get(j.t2Alias);
            if (n == null) { // 第一次遇到这个表
                n = new DefaultMutableTreeNode(j.t2Alias == null ? "子计划"
                    : (j.t2Alias + " (成本 = " + stats.get(table2Name).estimateScanCost() + ", 基数 = "
                       + stats.get(table2Name).estimateTableCardinality(selectivities.get(j.t2Alias)) + ")"));
                root.add(n);
            } else {
                // 添加已存在的右子节点
                root.add(n);
                neither = false;
            }
            m.put(j.t2Alias, root);

            // 如果该表与其他表连接
            if (!neither) {
                for (String key : m.keySet()) {
                    m.put(key, root); // 将所有表作为根的子节点
                }
            }

            treetop = root;
        }

        // 创建JTree组件并设置样式
        JTree tree = new JTree(treetop);
        JScrollPane treeView = new JScrollPane(tree);

        tree.setShowsRootHandles(true);

        // 设置叶子节点图标
        ImageIcon leafIcon = new ImageIcon("join.jpg");
        DefaultTreeCellRenderer renderer = new DefaultTreeCellRenderer();
        renderer.setOpenIcon(leafIcon);
        renderer.setClosedIcon(leafIcon);

        tree.setCellRenderer(renderer);

        f.setSize(300, 500);

        // 将树添加到窗口
        f.add(treeView);

        // 展开所有行以便查看
        for (int i = 0; i < tree.getRowCount(); i++) {
            tree.expandRow(i);
        }

        // 如果没有连接计划，显示提示信息
        if (js.size() == 0) {
            f.add(new JLabel("计划中没有连接。"));
        }

        f.pack();
    }
}
